﻿#include "Sm2.h"
#include <assert.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>


#define GM_GETU32(p)           \
    ((uint32_t) (p)[0] << 24 | \
     (uint32_t) (p)[1] << 16 | \
     (uint32_t) (p)[2] << 8 |  \
     (uint32_t) (p)[3])


#define GM_PUTU32(p, V)              \
    ((p)[0] = (uint8_t) ((V) >> 24), \
     (p)[1] = (uint8_t) ((V) >> 16), \
     (p)[2] = (uint8_t) ((V) >> 8),  \
     (p)[3] = (uint8_t) (V))

// SM2 P
const uint512_t GM_BN_P = {
        0xFFFFFFFF, 0xFFFFFFFF, 0x00000000, 0xFFFFFFFF,
        0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFE};

// SM2 N
const uint512_t GM_BN_N = {
        0x39D54123, 0x53BBF409, 0x21C6052B, 0x7203DF6B,
        0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFE};

// mont,(r ^ 2)(mod p)
static const uint512_t GM_BN_MONT_PRR = {
        0x00000003, 0x00000002, 0xFFFFFFFF, 0x00000002,
        0x00000001, 0x00000001, 0x00000002, 0x00000004};

// mont,(r ^ 2)(mod n)
static const uint512_t GM_BN_MONT_NRR = {
        0x7C114F20, 0x901192AF, 0xDE6FA2FA, 0x3464504A,
        0x3AFFE0D4, 0x620FC84C, 0xA22B3D3B, 0x1EB5E412};

// uint512_t(0)
static const uint512_t U512_0 = {
        0x0, 0x0, 0x0, 0x0,
        0x0, 0x0, 0x0, 0x0};

// uint512_t(1)
static const uint512_t U512_1 = {
        0x1, 0x0, 0x0, 0x0,
        0x0, 0x0, 0x0, 0x0};

// uint512_t(2)
static const uint512_t U512_2 = {
        0x2, 0x0, 0x0, 0x0,
        0x0, 0x0, 0x0, 0x0};

// mont 1 (mod p)
const uint512_t GM_BN_MONT_PONE = {
        0x00000001, 0x00000000, 0xFFFFFFFF, 0x00000000,
        0x00000000, 0x00000000, 0x00000000, 0x00000001};

// mont 1 (mod n)
static const uint512_t GM_BN_MONT_NONE = {
        0xC62ABEDD, 0xAC440BF6, 0xDE39FAD4, 0x8DFC2094,
        0x00000000, 0x00000000, 0x00000000, 0x00000001};

// mont SM2 G
static const u512_point_t GM_MONT_G = {
        {0xF418029E, 0x61328990, 0xDCA6C050, 0x3E7981ED, 0xAC24C3C3, 0xD6A1ED99, 0xE1C13B05, 0x91167A5E},
        {0x3C2D0DDD, 0xC1354E59, 0x8D3295FA, 0xC1F5E578, 0x6E2A48F8, 0x8D4CFB06, 0x81D735BD, 0x63CD65D4},
        {0x00000001, 0x00000000, 0xFFFFFFFF, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000001}};

const u512_point_t *GM_MONT_G_PTR = &GM_MONT_G;

// mont SM2 A
const uint512_t GM_BN_MONT_A = {
        0xFFFFFFFC, 0xFFFFFFFF, 0x00000003, 0xFFFFFFFC,
        0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFB};

// mont SM2 B
const uint512_t GM_BN_MONT_B = {
        0x2BC0DD42, 0x90D23063, 0xE9B537AB, 0x71CF379A,
        0x5EA51C3C, 0x52798150, 0xBA20E2C8, 0x240FE188};

// booth encoding
int u256_get_booth(const uint256_t a, int window_size, int i) {
    uint64_t mask = (1 << window_size) - 1;
    uint64_t wbits;
    int n, j;

    if (i == 0) {
        return (int) ((a[0] << 1) & mask) - (int) (a[0] & mask);
    }

    j = i * window_size - 1;
    n = j / 64;
    j = j % 64;

    wbits = a[n] >> j;
    if ((64 - j) < (int) (window_size + 1) && n < 3) {
        wbits |= a[n + 1] << (64 - j);
    }
    return (int) (wbits & mask) - (int) ((wbits >> 1) & mask);
}

// table[0, ..., 1 << (window_size - 1) - 1] = [1 * p, ..., (1 << (window_size - 1)) * p]
void make_point_for_window(int window_size, u512_point_t *table, const u512_point_t *p) {
    uint512_point_copy(&table[0], p);
    int i;
    for (i = 1; i < 1 << (window_size - 1); ++i) {
        if ((i + 1) % 2 == 0) {
            u512_point_double(&table[i], &table[(i + 1) / 2 - 1]);
        } else {
            u512_point_add(&table[i], &table[i - 1], p);
        }
    }
}

//#if defined(__GNUC__)
//#pragma GCC push_options
//#pragma GCC optimize("O3")
//#endif

// r[:8] = (a[:8] + b[:8])
static void uint512_i_bn_add_8(uint512_t r, const uint512_t a, const uint512_t b) {
    int i;
    r[0] = a[0] + b[0];
    for (i = 1; i < 8; i++) {
        r[i] = a[i] + b[i] + (r[i - 1] >> 32);
        r[i - 1] &= 0x0FFFFFFFFULL;
    }
}

// r[:11] = (a[:11] + b[:11])
static void uint512_i_bn_add_11(uint512_t r, const uint512_t a, const uint512_t b) {
    int i;
    r[0] = a[0] + b[0];
    for (i = 1; i < 11; i++) {
        r[i] = a[i] + b[i] + (r[i - 1] >> 32);
        r[i - 1] &= 0x0FFFFFFFFULL;
    }
}

// r = (a - b)
static void uint512_i_bn_sub(uint512_t r, const uint512_t a, const uint512_t b) {
    int i;
    r[0] = 0x100000000ULL + a[0] - b[0];
    for (i = 1; i < 7; i++) {
        r[i] = 0x0FFFFFFFFULL + a[i] - b[i] + (r[i - 1] >> 32);
        r[i - 1] &= 0x0FFFFFFFFULL;
    }
    r[i] = a[i] - b[i] + (r[i - 1] >> 32) - 1;
    r[i - 1] &= 0x0FFFFFFFFULL;
}

// r = k[b]
static void uint512_i_bn_mul(uint64_t *r, const uint64_t *k, const uint512_t b) {
    int i, j;
    // k0b0, k0b1+k1b0, .. ,k0b7+k1b6, k1b7
    uint64_t t;

    for (i = 0; i < 2; i++) {
        t = 0;
        for (j = 0; j < 8; j++) {
            t = r[i + j] + k[i] * b[j] + t;
            r[i + j] = t & 0x0FFFFFFFFULL;
            t >>= 32;
        }
        r[i + 8] = t;
    }
}

// r = (a + b)
static void uint512_i_bn_add(uint512_t r, const uint512_t a, const uint512_t b) {
    uint512_i_bn_add_8(r, a, b);
}

//#if defined(__GNUC__)
//#pragma GCC pop_options
//#endif

static int char_hex2int(char c) {
    switch (c) {
        case '0':
        case '1':
        case '2':
        case '3':
        case '4':
        case '5':
        case '6':
        case '7':
        case '8':
        case '9':
            return c - '0';
        case 'a':
        case 'b':
        case 'c':
        case 'd':
        case 'e':
        case 'f':
            return c - 'a' + 10;
        case 'A':
        case 'B':
        case 'C':
        case 'D':
        case 'E':
        case 'F':
            return c - 'A' + 10;
        default:
            return -1;
    }
}

int sm2_hex2bin(const char *in, int in_len, uint8_t *out) {
    int c;
    if ((in_len % 2) != 0) {
        return -1;
    }

    while (in_len) {
        if ((c = char_hex2int(*in++)) < 0) {
            return -1;
        }
        *out = (uint8_t) c << 4;

        if ((c = char_hex2int(*in++)) < 0) {
            return -1;
        }
        *out |= (uint8_t) c;

        in_len -= 2;
        out++;
    }
    return 1;
}

void uint512_to_bytes(const uint512_t a, uint8_t out[32]) {
    int i;
    for (i = 7; i >= 0; i--) {
        GM_PUTU32(out, a[i]);
        out += 4;
    }
}

void uint512_from_bytes(uint512_t r, const uint8_t in[32]) {
    int i;
    for (i = 7; i >= 0; i--) {
        r[i] = GM_GETU32(in);
        in += 4;
    }
}

void uint512_to_hex(const uint512_t a, char hex[64]) {
    int i;
    for (i = 7; i >= 0; i--) {
        int len;
        len = sprintf(hex, "%08X", (uint32_t) a[i]);
        assert(len == 8);
        hex += 8;
    }
}

int uint512_from_hex(uint512_t r, const char hex[64]) {
    uint8_t buf[32];
    if (sm2_hex2bin(hex, 64, buf) < 0)
        return -1;
    uint512_from_bytes(r, buf);
    return 1;
}

int uint512_cmp(const uint512_t a, const uint512_t b) {
    for (int i = 7; i >= 0; i--) {
        if (a[i] > b[i])
            return 1;
        if (a[i] < b[i])
            return -1;
    }
    return 0;
}

// r = (a + b)(mod m) precondition: a < m && b < m
void uint512_add(uint512_t r, const uint512_t a, const uint512_t b, const uint512_t m) {
    uint512_i_bn_add(r, a, b);
    if (uint512_cmp(r, m) >= 0) {
        uint512_i_bn_sub(r, r, m);
    }
}

// r = (a - b)(mod m)
void uint512_sub(uint512_t r, const uint512_t a, const uint512_t b, const uint512_t m) {
    if (uint512_cmp(a, b) >= 0) {
        uint512_i_bn_sub(r, a, b);
    } else {
        uint512_t t;
        uint512_i_bn_sub(t, m, b);
        uint512_i_bn_add(r, t, a);
    }
}

// r = (a * b)(mod m)
void uint512_mont_mul(uint512_t r, const uint512_t a, const uint512_t b, const uint512_t m) {
    uint64_t temp_r[11] = {0};
    uint64_t temp_mul[11] = {0};
    uint64_t temp_q[2] = {0};
    int i;

    uint64_t mod_t = (m[0] == GM_BN_N[0] ? 0x327F9E8872350975ULL : 0x01);

    for (i = 0; i < 4; i++) {
        // bi*a
        memset(temp_mul, 0, sizeof(uint64_t) * 11);
        uint512_i_bn_mul(temp_mul, &b[i * 2], a);
        uint512_i_bn_add_11(temp_r, temp_r, temp_mul);

        // cal q
        uint64_t q = (temp_r[0] | temp_r[1] << 32) * mod_t;
        temp_q[0] = q & 0x0FFFFFFFFULL;
        temp_q[1] = q >> 32;

        // q * n
        memset(temp_mul, 0, sizeof(uint64_t) * 11);
        uint512_i_bn_mul(temp_mul, temp_q, m);
        uint512_i_bn_add_11(temp_r, temp_r, temp_mul);

        // reduce temp_r
        memmove(temp_r, temp_r + 2, sizeof(uint64_t) * 9);
        temp_r[9] = 0;
        temp_r[10] = 0;
    }

    temp_r[7] += temp_r[8] << 32;
    if (uint512_cmp(temp_r, m) >= 0) {
        uint512_i_bn_sub(temp_r, temp_r, m);
    }
    uint512_copy(r, temp_r);
}

void uint512_to_mont(uint512_t r, const uint512_t a, const uint512_t m) {
    uint512_mont_mul(r, a, (m[0] == GM_BN_N[0] ? GM_BN_MONT_NRR : GM_BN_MONT_PRR), m);
}

void uint512_from_mont(uint512_t r, const uint512_t a, const uint512_t m) {
    uint512_mont_mul(r, a, U512_1, m);
}

void uint512_sqr(uint512_t r, const uint512_t a, const uint512_t m) {
    uint512_mont_mul(r, a, a, m);
}

void uint512_exp(uint512_t r, const uint512_t a, const uint512_t b, const uint512_t m) {
    uint512_t t;
    uint64_t w;
    int i, j;

    // set t to mont one
    uint512_to_mont(t, U512_1, m);

    for (i = 7; i >= 0; i--) {
        w = b[i];
        for (j = 0; j < 32; j++) {
            uint512_sqr(t, t, m);
            if (w & 0x080000000ULL) {
                uint512_mont_mul(t, t, a, m);
            }
            w <<= 1;
        }
    }

    uint512_copy(r, t);
}

void uint512_inv(uint512_t r, const uint512_t a, const uint512_t m) {
    uint512_t e;
    uint512_i_bn_sub(e, m, U512_2);
    uint512_exp(r, a, e, m);
}

void uint512_set_mont_one(uint512_t r) {
    uint512_copy(r, GM_BN_MONT_PONE);
}

void uint512_set_zero(uint512_t r) {
    uint512_copy(r, U512_0);
}

int uint512_is_mont_one(const uint512_t r) {
    return memcmp(r, GM_BN_MONT_PONE, sizeof(uint512_t)) == 0;
}

int uint512_is_zero(const uint512_t r) {
    return uint512_cmp(r, U512_0) == 0;
}

void uint512_point_init(u512_point_t *r) {
    memset(r, 0, sizeof(u512_point_t));
    uint512_set_mont_one(r->X);
    uint512_set_mont_one(r->Y);
}

void uint512_point_set_xy(u512_point_t *r, const uint512_t x, const uint512_t y) {
    uint512_to_mont(r->X, x, GM_BN_P);
    uint512_to_mont(r->Y, y, GM_BN_P);
    uint512_set_mont_one(r->Z);
}

void u512_point_get_xy_mont(const u512_point_t *p, uint512_t x, uint512_t y) {
    uint512_t z_inv;
    if (uint512_is_mont_one(p->Z)) {
        if (x) {
            uint512_copy(x, p->X);
        }
        if (y) {
            uint512_copy(y, p->Y);
        }
    } else {
        // z^{-1}
        uint512_inv(z_inv, p->Z, GM_BN_P);
        if (y) {
            uint512_mont_mul(y, p->Y, z_inv, GM_BN_P);
        }
        // z^{-2}
        uint512_sqr(z_inv, z_inv, GM_BN_P);
        if (x) {
            uint512_mont_mul(x, p->X, z_inv, GM_BN_P);
        }
        if (y) {
            uint512_mont_mul(y, y, z_inv, GM_BN_P);
        }
    }
}

void uint512_point_get_xy(const u512_point_t *p, uint512_t x, uint512_t y) {
    u512_point_get_xy_mont(p, x, y);
    if (x) {
        uint512_from_mont(x, x, GM_BN_P);
    }
    if (y) {
        uint512_from_mont(y, y, GM_BN_P);
    }
}

void u256_point_to_u512_point(int negative, u512_point_t *u512_point, const u256_point_t *u256_point) {
    u256_to_u512(&u512_point->X, &(u256_point->X));
    u256_to_u512(&u512_point->Y, &(u256_point->Y));
    if (negative) {
        uint512_i_bn_sub(u512_point->Y, GM_BN_P, u512_point->Y);
    }
    uint512_set_mont_one(u512_point->Z);
}

int u512_point_is_at_infinity(const u512_point_t *p) {
    return uint512_is_zero(p->Z);
}

void u512_point_from_bytes(u512_point_t *r, const uint8_t in[64]) {
    uint512_t x;
    uint512_t y;
    uint512_from_bytes(x, in);
    uint512_from_bytes(y, in + 32);
    uint512_point_set_xy(r, x, y);
}

void u512_point_double(u512_point_t *r, const u512_point_t *p) {
    uint512_t tmp1, tmp2;
    uint512_t X3, Y3, Z3;

    uint512_t R1, R2, R3;

    if (u512_point_is_at_infinity(p)) {
        uint512_point_copy(r, p);
        return;
    }

    // λ1 = 3X^2 + aZ^4 = 3X^2 - 3Z^4 = 3 * ( (X + Z^2) * (X - Z^2) )
    // z^2
    uint512_sqr(tmp1, p->Z, GM_BN_P);
    // X + Z^2
    uint512_add(tmp2, p->X, tmp1, GM_BN_P);
    // X - Z^2
    uint512_sub(tmp1, p->X, tmp1, GM_BN_P);
    // (X + Z^2) * (X - Z^2)
    uint512_mont_mul(R1, tmp1, tmp2, GM_BN_P);
    // λ1 = 3 * (X + Z^2) * (X - Z^2)
    uint512_add(tmp1, R1, R1, GM_BN_P);
    uint512_add(R1, tmp1, R1, GM_BN_P);

    // λ2 = X4Y^2 = X * (2Y)^2
    // Z3 = 2YZ

    // 2Y
    uint512_add(tmp1, p->Y, p->Y, GM_BN_P);
    // Z3 = 2YZ
    uint512_mont_mul(Z3, tmp1, p->Z, GM_BN_P);

    // λ2 = X * (2Y)^2
    uint512_sqr(R2, tmp1, GM_BN_P);

    uint512_mont_mul(R2, p->X, R2, GM_BN_P);

    // λ3 = 8Y^4 = 2 * 4Y^4 = 2 * (2 * Y^2) ^ 2
    // Y^2
    uint512_sqr(tmp1, p->Y, GM_BN_P);
    // 2Y^2
    uint512_add(tmp1, tmp1, tmp1, GM_BN_P);
    // (2 * Y^2) ^ 2
    uint512_sqr(tmp1, tmp1, GM_BN_P);
    // λ3 = 2 * (2 * Y^2) ^ 2
    uint512_add(R3, tmp1, tmp1, GM_BN_P);

    // X3 = λ1^2 − 2λ2
    // R1^2
    uint512_sqr(tmp1, R1, GM_BN_P);
    // 2R2
    uint512_add(tmp2, R2, R2, GM_BN_P);
    // X3 = λ1^2 − 2λ2
    uint512_sub(X3, tmp1, tmp2, GM_BN_P);

    // Y3 = λ1(λ2 − X3) − λ3
    // λ2 − X3
    uint512_sub(tmp1, R2, X3, GM_BN_P);
    // λ1(λ2 − X3)
    uint512_mont_mul(tmp2, R1, tmp1, GM_BN_P);
    // Y3 = λ1(λ2 − X3) − λ3
    uint512_sub(Y3, tmp2, R3, GM_BN_P);

    // output
    uint512_copy(r->X, X3);
    uint512_copy(r->Y, Y3);
    uint512_copy(r->Z, Z3);
}

void u512_point_add(u512_point_t *r, const u512_point_t *a, const u512_point_t *b) {
    // U1 = X1 * (Z2)^2
    // U2 = X2 * (Z1)^2
    // S1 = Y1 * (Z2)^3
    // S2 = Y2 * (Z1)^3
    // if U1==U2 && S1==S2 return double(a)
    // H = U2 - U1
    // R = S2 - S1
    // X3 = R^2 - H^3 - 2U1 * H^2
    // Y3 = R * ( U1 * H^2 - X3) - S1 * H^3
    // Z3 = H * Z1 * Z2

    uint512_t U1, U2;
    uint512_t S1, S2;
    uint512_t H;
    uint512_t R;

    uint512_t X3, Y3, Z3;

    if (u512_point_is_at_infinity(a)) {
        uint512_point_copy(r, b);
        return;
    }

    if (u512_point_is_at_infinity(b)) {
        uint512_point_copy(r, a);
        return;
    }

    // Z1 ^ 2
    uint512_sqr(H, a->Z, GM_BN_P);
    // Z2 ^ 2
    uint512_sqr(R, b->Z, GM_BN_P);
    // U1 = X1 * (Z2)^2
    uint512_mont_mul(U1, a->X, R, GM_BN_P);
    // U2 = X2 * (Z1)^2
    uint512_mont_mul(U2, b->X, H, GM_BN_P);
    // Z1 ^ 3
    uint512_mont_mul(H, H, a->Z, GM_BN_P);
    // Z2 ^ 3
    uint512_mont_mul(R, R, b->Z, GM_BN_P);
    // S1 = Y1 * (Z2)^3
    uint512_mont_mul(S1, a->Y, R, GM_BN_P);
    // S2 = Y2 * (Z1)^3
    uint512_mont_mul(S2, b->Y, H, GM_BN_P);

    if (uint512_cmp(U1, U2) == 0 && uint512_cmp(S1, S2) == 0) {
        u512_point_double(r, a);
        return;
    }

    // H = U2 - U1
    uint512_sub(H, U2, U1, GM_BN_P);
    // R = S2 - S1
    uint512_sub(R, S2, S1, GM_BN_P);
    // Z3 = H * Z1 * Z2
    uint512_mont_mul(Z3, H, a->Z, GM_BN_P);
    uint512_mont_mul(Z3, Z3, b->Z, GM_BN_P);
    // H ^ 2
    uint512_sqr(U2, H, GM_BN_P);
    // H ^ 3
    uint512_mont_mul(H, U2, H, GM_BN_P);
    // U1 * H ^2
    uint512_mont_mul(Y3, U1, U2, GM_BN_P);
    // 2 * U1 * H ^2
    uint512_add(U2, Y3, Y3, GM_BN_P);
    // R^2
    uint512_sqr(X3, R, GM_BN_P);
    // X3 = R^2 - H^3 - 2U1 * H^2
    uint512_sub(X3, X3, H, GM_BN_P);
    uint512_sub(X3, X3, U2, GM_BN_P);
    // Y3 = R * ( U1 * H^2 - X3) - S1 * H^3
    // ( U1 * H^2 - X3)
    uint512_sub(Y3, Y3, X3, GM_BN_P);
    // R * ( U1 * H^2 - X3)
    uint512_mont_mul(Y3, Y3, R, GM_BN_P);
    // S1 * H^3
    uint512_mont_mul(S1, S1, H, GM_BN_P);
    // Y3 = R * ( U1 * H^2 - X3) - S1 * H^3
    uint512_sub(Y3, Y3, S1, GM_BN_P);

    // output
    uint512_copy(r->X, X3);
    uint512_copy(r->Y, Y3);
    uint512_copy(r->Z, Z3);
}

void u512_point_mul(u512_point_t *r, const uint512_t *k, const u512_point_t *p) {
    u512_point_t tmp_table[16];
    int window_size = 5;
    int is_infinity = 1;
    int n = (256 + window_size - 1) / window_size;
    int i;
    uint256_t u256_k;
    u512_to_u256(&u256_k, k);
    make_point_for_window(window_size, tmp_table, p);
    u512_point_t tmp_point;
    for (i = n - 1; i >= 0; i--) {
        int booth = u256_get_booth(u256_k, window_size, i);
        if (is_infinity) {
            if (booth != 0) {
                uint512_point_copy(r, &tmp_table[booth - 1]);
                is_infinity = 0;
            }
        } else {
            for (int j = 0; j < window_size; ++j) {
                u512_point_double(r, r);
            }
            if (booth != 0) {
                uint512_point_copy(&tmp_point, &tmp_table[abs(booth) - 1]);
                if (booth < 0) {
                    uint512_i_bn_sub(tmp_point.Y, GM_BN_P, tmp_point.Y);
                }
                u512_point_add(r, r, &tmp_point);
            }
        }
    }
}

extern const u256_point_t table_list[2368];

void u512_point_mul_G(u512_point_t *r, const uint512_t *k) {
    int window_size = 7;
    int batch = (256 + window_size - 1) / window_size;
    int batch_num = 1 << (window_size - 1);

    u512_point_t tmp_point;
    uint256_t u256_k;
    u512_to_u256(&u256_k, k);

    int is_infinity = 1;
    int i;

    for (i = batch - 1; i >= 0; i--) {
        int booth = u256_get_booth(u256_k, window_size, i);
        if (is_infinity) {
            if (booth != 0) {
                u256_point_to_u512_point(0, r, &table_list[i * batch_num + booth - 1]);
                is_infinity = 0;
            }
        } else {
            if (0 == booth) {
                u512_point_add(&tmp_point, r, GM_MONT_G_PTR);
            } else {
                u256_point_to_u512_point(booth < 0, &tmp_point, &table_list[i * batch_num + abs(booth) - 1]);
                u512_point_add(r, r, &tmp_point);
            }
        }
    }
}

void u512_point_decode(u512_point_t *p, const uint8_t *in) {
    if (in[0] == 0x04) {
        u512_point_from_bytes(p, in + 1);
    } else if (in[0] == 0x02 || in[0] == 0x03) {
        int yTile = in[0] - 0x02;
        uint512_t x, r;

        // (P + 1) / 4
        uint512_t padd1shit2 = {
                0x00000000, 0x40000000, 0xC0000000, 0xFFFFFFFF,
                0xFFFFFFFF, 0xFFFFFFFF, 0xBFFFFFFF, 0x3FFFFFFF};

        uint512_from_bytes(x, in + 1);
        uint512_to_mont(x, x, GM_BN_P);

        // r = x ^ 2
        uint512_sqr(r, x, GM_BN_P);

        // r = x^2 + a
        uint512_add(r, r, GM_BN_MONT_A, GM_BN_P);

        // r = x^3 + ax
        uint512_mont_mul(r, r, x, GM_BN_P);

        // r = x^3 + ax + b
        uint512_add(r, r, GM_BN_MONT_B, GM_BN_P);

        // r = sqrt(x^3 + ax + b) = (x^3 + ax + b) ^ ((P + 1) / 4) (mod P)
        uint512_exp(r, r, padd1shit2, GM_BN_P);

        uint512_from_mont(padd1shit2, r, GM_BN_P);
        if ((padd1shit2[0] & 0x01) != yTile) {
            uint512_sub(r, GM_BN_P, padd1shit2, GM_BN_P);
            uint512_to_mont(r, r, GM_BN_P);
        }

        uint512_copy(p->X, x);
        uint512_copy(p->Y, r);
        uint512_set_mont_one(p->Z);
    }else{
        u512_point_from_bytes(p, in);
    }
}

int sm2_do_sign(const uint512_t key, const uint512_t digest, unsigned char sig[64]) {
    u512_point_t _P, *P = &_P;
    uint512_t d;
    uint512_t e;
    uint512_t k;
    uint512_t x;
    uint512_t r;
    uint512_t s;

    if (!key || !digest || !sig) {
        return -1;
    }

    uint512_to_mont(d, key, GM_BN_N);

    // e = H(M)
    uint512_copy(e, digest);
    uint8_t buf[32];
retry:
    // rand k in [1, n - 1]
    do {
        do {
            srand(time(0));
            for (int a = 0; a < 32; a++) {
                buf[a] = (uint8_t) (rand() % 255);
            }
            uint512_from_bytes(k, buf);
        } while (uint512_cmp(k, GM_BN_N) >= 0);
    } while (uint512_is_zero(k));
    uint512_copy(k, key);
    // (x, y) = kG
    u512_point_mul_G(P, &k);
    uint512_point_get_xy(P, x, NULL);


    // r = e + x (mod n)
    uint512_add(r, e, x, GM_BN_N);

    /* if r == 0 or r + k == n re-generate k */
    if (uint512_is_zero(r)) {
        goto retry;
    }
    uint512_add(x, r, k, GM_BN_N);
    if (uint512_is_zero(x)) {
        goto retry;
    }

    uint512_to_bytes(r, sig);

    /* s = ((1 + d)^-1 * (k - r * d)) mod n */
    uint512_to_mont(r, r, GM_BN_N);
    uint512_to_mont(k, k, GM_BN_N);

    uint512_mont_mul(e, r, d, GM_BN_N);
    uint512_sub(k, k, e, GM_BN_N);

    uint512_add(x, GM_BN_MONT_NONE, d, GM_BN_N);
    uint512_inv(x, x, GM_BN_N);
    uint512_mont_mul(s, x, k, GM_BN_N);

    if (uint512_is_zero(s)) {
        goto retry;
    }

    uint512_from_mont(s, s, GM_BN_N);
    uint512_to_bytes(s, sig + 32);
    return 1;
}

int sm2_do_verify(const u512_point_t *key, const uint512_t digest, const unsigned char sig[64]) {
    u512_point_t _P, *P = &_P;
    u512_point_t _Q, *Q = &_Q;
    uint512_t r;
    uint512_t s;
    uint512_t e;
    uint512_t x;
    uint512_t t;

    if (!key || !digest || !sig) {
        return -1;
    }

    // parse signature values
    uint512_from_bytes(r, (const uint8_t *) sig);
    uint512_from_bytes(s, (const uint8_t *) sig + 32);
    if (uint512_is_zero(r) || uint512_cmp(r, GM_BN_N) >= 0 || uint512_is_zero(s) || uint512_cmp(s, GM_BN_N) >= 0) {
        return -1;
    }

    // parse public key
    uint512_point_copy(P, key);

    // t = r + s (mod n)
    // check t != 0
    uint512_add(t, r, s, GM_BN_N);
    if (uint512_is_zero(t)) {
        return -1;
    }

    // Q = s * G + t * P
    u512_point_mul_G(Q, &s);
    u512_point_mul(P, &t, P);
    u512_point_add(Q, Q, P);
    uint512_point_get_xy(Q, x, NULL);

    // e  = H(M)
    // r' = e + x (mod n)
    uint512_copy(e, digest);
    uint512_add(e, e, x, GM_BN_N);

    // check if r == r'
    if (uint512_cmp(e, r) == 0) {
        return 1;
    } else {
        return 0;
    }
}

// SM2 a
static const unsigned char GM_ECC_A[] = {
        0xFF, 0xFF, 0xFF, 0xFE, 0xFF, 0xFF, 0xFF, 0xFF,
        0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
        0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x00,
        0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFC};

// SM2 b
static const unsigned char GM_ECC_B[] = {
        0x28, 0xE9, 0xFA, 0x9E, 0x9D, 0x9F, 0x5E, 0x34,
        0x4D, 0x5A, 0x9E, 0x4B, 0xCF, 0x65, 0x09, 0xA7,
        0xF3, 0x97, 0x89, 0xF5, 0x15, 0xAB, 0x8F, 0x92,
        0xDD, 0xBC, 0xBD, 0x41, 0x4D, 0x94, 0x0E, 0x93};

// SM2 Gx
static const unsigned char GM_ECC_G_X[] = {
        0x32, 0xC4, 0xAE, 0x2C, 0x1F, 0x19, 0x81, 0x19,
        0x5F, 0x99, 0x04, 0x46, 0x6A, 0x39, 0xC9, 0x94,
        0x8F, 0xE3, 0x0B, 0xBF, 0xF2, 0x66, 0x0B, 0xE1,
        0x71, 0x5A, 0x45, 0x89, 0x33, 0x4C, 0x74, 0xC7};

// SM2 Gy
static const unsigned char GM_ECC_G_Y[] = {
        0xBC, 0x37, 0x36, 0xA2, 0xF4, 0xF6, 0x77, 0x9C,
        0x59, 0xBD, 0xCE, 0xE3, 0x6B, 0x69, 0x21, 0x53,
        0xD0, 0xA9, 0x87, 0x7C, 0xC6, 0x2A, 0x47, 0x40,
        0x02, 0xDF, 0x32, 0xE5, 0x21, 0x39, 0xF0, 0xA0};

// SM2 N-1
static const uint512_t GM_BN_N_SUB_ONE = {
        0x39D54122, 0x53BBF409, 0x21C6052B, 0x7203DF6B,
        0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFE};

static const uint512_t GM_BN_2W = {
        0x00000000, 0x00000000, 0x00000000, 0x80000000,
        0x00000000, 0x00000000, 0x00000000, 0x00000000};

static const uint512_t GM_BN_2W_SUB_ONE = {
        0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0x7FFFFFFF,
        0x00000000, 0x00000000, 0x00000000, 0x00000000};


#ifndef GM_PUT_UINT32_BE
#define GM_PUT_UINT32_BE(n, b, i)                   \
    {                                               \
        (b)[(i)] = (unsigned char) ((n) >> 24);     \
        (b)[(i) + 1] = (unsigned char) ((n) >> 16); \
        (b)[(i) + 2] = (unsigned char) ((n) >> 8);  \
        (b)[(i) + 3] = (unsigned char) ((n));       \
    }
#endif

void sm2_compute_z_digest(const unsigned char *id_bytes, unsigned int idLen, const u512_point_t *pub_key,
                          unsigned char output[32]) {
    gm_sm3_context _ctx;
    gm_sm3_context *ctx = &_ctx;
    uint512_t x, y;

    gm_sm3_init(ctx);

    // update idlen
    ctx->buf[0] = ((idLen >> 5) & 0x0FF);
    ctx->buf[1] = ((idLen << 3) & 0x0FF);
    ctx->cur_buf_len = 2;

    // update id
    gm_sm3_update(ctx, id_bytes, idLen);

    // update a
    gm_sm3_update(ctx, GM_ECC_A, 32);

    // update b
    gm_sm3_update(ctx, GM_ECC_B, 32);

    // update Gx
    gm_sm3_update(ctx, GM_ECC_G_X, 32);

    // update Gy
    gm_sm3_update(ctx, GM_ECC_G_Y, 32);

    uint512_point_get_xy(pub_key, x, y);

    // update Px
    uint512_to_bytes(x, output);// 借用output当缓冲区
    gm_sm3_update(ctx, output, 32);

    // update Py
    uint512_to_bytes(y, output);// 借用output当缓冲区
    gm_sm3_update(ctx, output, 32);

    gm_sm3_done(ctx, output);
}

static int sm2_check_public_key(const u512_point_t *pub_key) {
    if (u512_point_is_at_infinity(pub_key)) {
        return 0;
    }
    uint512_t x, y, r;
    uint512_point_get_xy(pub_key, x, y);

    if (uint512_is_zero(x) || uint512_cmp(x, GM_BN_P) >= 0) {
        return 0;
    }

    if (uint512_is_zero(y) || uint512_cmp(y, GM_BN_P) >= 0) {
        return 0;
    }

    //y^2 = x^3 + ax + b
    uint512_to_mont(x, x, GM_BN_P);
    uint512_to_mont(y, y, GM_BN_P);

    // r = x ^ 2
    uint512_sqr(r, x, GM_BN_P);

    // r = x^2 + a
    uint512_add(r, r, GM_BN_MONT_A, GM_BN_P);

    // r = x^3 + ax
    uint512_mont_mul(r, r, x, GM_BN_P);

    // r = x^3 + ax + b
    uint512_add(r, r, GM_BN_MONT_B, GM_BN_P);

    uint512_sqr(y, y, GM_BN_P);

    if (uint512_cmp(r, y) != 0) {
        return 0;
    }

    return 1;
}

int recover_private_key(sm2_context *ctx, const unsigned char *key, unsigned int kLen) {
    if (kLen != 32) {
        return 0;
    }
    uint512_from_bytes(ctx->private_key, key);
    // check k ∈ [1, n-2]
    if (uint512_is_zero(ctx->private_key) || uint512_cmp(ctx->private_key, GM_BN_N_SUB_ONE) >= 0) {
        return 0;
    }
    // check public key
    u512_point_mul_G(&ctx->public_key, &ctx->private_key);
    if (sm2_check_public_key(&ctx->public_key) != 1) {
        return 0;
    }
    return 1;
}

static int recover_public_key(sm2_context *ctx, const unsigned char *key, unsigned int kLen) {
	if ((kLen != 33 && kLen != 64 && kLen != 65)
			|| (kLen != 64
					&& (key[0] != 0x04 && key[0] != 0x02 && key[0] != 0x03))) {
		return 0;
	}
    // check public key
    u512_point_decode(&ctx->public_key, key);
    if (sm2_check_public_key(&ctx->public_key) != 1) {
        return 0;
    }
    return 1;
}

int sm2_sign_init(sm2_context *ctx, const unsigned char *key, unsigned int kLen,
                  const unsigned char *id_bytes, unsigned int idLen, int forSign) {
    if (forSign) {
        // private
        if (recover_private_key(ctx, key, kLen) == 0) {
            return 0;
        }
    } else {
        // public
        if (recover_public_key(ctx, key, kLen) == 0) {
            return 0;
        }
    }
    // compute z digest
    sm2_compute_z_digest(id_bytes, idLen, &ctx->public_key, ctx->buf);
    gm_sm3_init(&ctx->sm3_ctx);
    gm_sm3_update(&ctx->sm3_ctx, ctx->buf, 32);
    ctx->state = forSign;

    return 1;
}

void sm2_sign_update(sm2_context *ctx, const unsigned char *input, unsigned int iLen) {
    gm_sm3_update(&ctx->sm3_ctx, input, iLen);
}

int sm2_sign_verify(sm2_context *ctx, unsigned char sig[64]) {
    if (ctx->state) {
        return 0;
    }
    uint512_t digest;

    gm_sm3_done(&ctx->sm3_ctx, ctx->buf);

    uint512_from_bytes(digest, ctx->buf);
    return sm2_do_verify(&ctx->public_key, digest, sig);
}

int sm2_sign_generate(sm2_context *ctx, unsigned char sig[64]) {
    if (!ctx->state) {
        return 0;
    }
    uint512_t digest;

    gm_sm3_done(&ctx->sm3_ctx, ctx->buf);

    uint512_from_bytes(digest, ctx->buf);

    return sm2_do_sign(ctx->private_key, digest, sig);
}

int sm2_crypt_init(sm2_context *ctx, const unsigned char *key, unsigned int kLen, int forEncryption, unsigned char c1[64]) {
    uint512_t k;
    uint8_t buf[256] = {1, 2, 3, 4, 5, 6, 7, 8};
    if (forEncryption) {
        // rand k in [1, n - 1]
        do {
            do {
                uint512_from_bytes(k, buf);
            } while (uint512_cmp(k, GM_BN_N) >= 0);
        } while (uint512_is_zero(k));
    }
    return sm2_crypt_init_ex(ctx, key, kLen, forEncryption, c1, &k);
}

void u512_point_to_bytes(const u512_point_t *p, uint8_t out[64]) {
    uint512_t x;
    uint512_t y;
    uint512_point_get_xy(p, x, y);
    uint512_to_bytes(x, out);
    uint512_to_bytes(y, out + 32);
}

int sm2_crypt_init_ex(sm2_context *ctx, const unsigned char *key, unsigned int kLen,
                   int forEncryption, unsigned char c1[64], const uint512_t *test_key) {
    u512_point_t p;
    if (forEncryption) {
        if (recover_public_key(ctx, key, kLen) == 0) {
            return 0;
        }
        u512_point_mul(&p, test_key, &GM_MONT_G);
        u512_point_to_bytes(&p, c1);
        u512_point_mul(&p, test_key, &ctx->public_key);
        u512_point_to_bytes(&p, ctx->x2y2);
    } else {
        if (recover_private_key(ctx, key, kLen) == 0) {
            return 0;
        }
        u512_point_from_bytes(&p, c1);
        u512_point_mul(&p, &ctx->private_key, &p);
        u512_point_to_bytes(&p, ctx->x2y2);
    }
    ctx->state = forEncryption;
    ctx->cur_buf_len = 0;
    ctx->ct = 1;
    gm_sm3_init(&ctx->sm3_ctx);
    gm_sm3_update(&ctx->sm3_ctx, ctx->x2y2, 32);
    return 1;
}

static void crypt_update_one_round(sm2_context *ctx, unsigned char *output, int len) {
    int i;
    // KDF
    gm_sm3_context sm3_ctx;
    gm_sm3_init(&sm3_ctx);
    gm_sm3_update(&sm3_ctx, ctx->x2y2, 64);

    GM_PUT_UINT32_BE(ctx->ct, output, 0);
    gm_sm3_update(&sm3_ctx, output, 4);
    gm_sm3_done(&sm3_ctx, output);

    for (i = 0; i < len; i++) {
        output[i] ^= ctx->buf[i];
    }

    if (ctx->state) {
        gm_sm3_update(&ctx->sm3_ctx, ctx->buf, len);
    } else {
        gm_sm3_update(&ctx->sm3_ctx, output, len);
    }

    ctx->cur_buf_len = 0;
    ctx->ct++;
}

int sm2_crypt_update(sm2_context *ctx, const unsigned char *input, unsigned int iLen, unsigned char *output) {
    int rLen = 0;

    while (iLen--) {
        ctx->buf[ctx->cur_buf_len++] = *input++;
        if (ctx->cur_buf_len == 32) { // is full
            crypt_update_one_round(ctx, output + rLen, 32);
            rLen += 32;
        }
    }
    return rLen;
}

int sm2_crypt_done(sm2_context *ctx, unsigned char *output, unsigned char c3[32]) {
    int rLen = ctx->cur_buf_len;
    crypt_update_one_round(ctx, output, rLen);

    gm_sm3_update(&ctx->sm3_ctx, ctx->x2y2 + 32, 32);
    gm_sm3_done(&ctx->sm3_ctx, c3);
    return rLen;
}
